/* http://www.gnu.org/copyleft/gpl.html */
// Code by gynvael.coldwind//vx
// !!!Use at your own risk!!!
#include<stdio.h>
#include<string.h>
#include<windows.h>
#include<winnt.h>

_CRTIMP int __cdecl	stricmp (const char*, const char*);

static IMAGE_IMPORT_DESCRIPTOR *FindImportDescriptor(const char *DllName)
{
  unsigned long MainVA = (unsigned long)GetModuleHandle(0);

  if(*(unsigned short*)MainVA != *(unsigned short*)"MZ") return NULL;

  IMAGE_DOS_HEADER *DosHeader = (IMAGE_DOS_HEADER*)MainVA;
  IMAGE_NT_HEADERS *NtHeaders = (IMAGE_NT_HEADERS*)(MainVA + DosHeader->e_lfanew);

  if(NtHeaders->Signature != IMAGE_NT_SIGNATURE) return NULL;

  IMAGE_DATA_DIRECTORY *ImportDirectory = &NtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
  IMAGE_IMPORT_DESCRIPTOR *ImportDescriptor = (IMAGE_IMPORT_DESCRIPTOR*)(MainVA + ImportDirectory->VirtualAddress);

  for(; ImportDescriptor->Name; ImportDescriptor++)
    if(stricmp((const char*)(MainVA + ImportDescriptor->Name), DllName) == 0) break;

  if(ImportDescriptor->Name == 0) return NULL;

  return ImportDescriptor;
}

void* DllSpoof_SpoofImportByOrd(const char *DllName, unsigned short Ord, void *SpoofFuncAddress)
{
  unsigned long MainVA = (unsigned long)GetModuleHandle(0);

  IMAGE_IMPORT_DESCRIPTOR *ImportDescriptor = FindImportDescriptor(DllName);
  if(ImportDescriptor == NULL) return NULL;

  // Seek and spoof
  IMAGE_THUNK_DATA32 *FThunk, *OThunk;
  FThunk = (IMAGE_THUNK_DATA32*)(MainVA + ImportDescriptor->FirstThunk);
  OThunk = (IMAGE_THUNK_DATA32*)(MainVA + ImportDescriptor->OriginalFirstThunk);

  void *OriginalFuncAddr;
  for(; OThunk->u1.AddressOfData; OThunk++, FThunk++)
  {
    if(!(OThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG32)) continue;
    if((unsigned short)OThunk->u1.Ordinal != Ord) continue;
    OriginalFuncAddr = (void*)FThunk->u1.Function;
    DWORD OldRight;
    VirtualProtect(&FThunk->u1.Function, 4, PAGE_EXECUTE_READWRITE, &OldRight);
    FThunk->u1.Function = (DWORD)SpoofFuncAddress;
    DWORD Temp;
    VirtualProtect(&FThunk->u1.Function, 4, OldRight, &Temp);
    break;
  }

  return OriginalFuncAddr;
}

void* DllSpoof_SpoofImport(const char *DllName, const char *FuncName, void *SpoofFuncAddress)
{
  unsigned long MainVA = (unsigned long)GetModuleHandle(0);

  IMAGE_IMPORT_DESCRIPTOR *ImportDescriptor = FindImportDescriptor(DllName);
  if(ImportDescriptor == NULL) return NULL;

  // Seek and spoof
  IMAGE_THUNK_DATA32 *FThunk, *OThunk;
  FThunk = (IMAGE_THUNK_DATA32*)(MainVA + ImportDescriptor->FirstThunk);
  OThunk = (IMAGE_THUNK_DATA32*)(MainVA + ImportDescriptor->OriginalFirstThunk);

  void *OriginalFuncAddr;
  for(; OThunk->u1.AddressOfData; OThunk++, FThunk++)
  {
    if(OThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG32) continue;
    IMAGE_IMPORT_BY_NAME *ImportFuncName = (IMAGE_IMPORT_BY_NAME*)(MainVA + OThunk->u1.AddressOfData);
     if(strcmp((const char*)ImportFuncName->Name, FuncName) != 0) continue;
    OriginalFuncAddr = (void*)FThunk->u1.Function;

    DWORD OldRight;
    VirtualProtect(&FThunk->u1.Function, 4, PAGE_EXECUTE_READWRITE, &OldRight);
    FThunk->u1.Function = (DWORD)SpoofFuncAddress;
    DWORD Temp;
    VirtualProtect(&FThunk->u1.Function, 4, OldRight, &Temp);
    break;
  }

  return OriginalFuncAddr;
}

BOOL DllSpoof_RecreateImports(const char *DllName, HINSTANCE hOriginal)
{
  unsigned long MainVA = (unsigned long)GetModuleHandle(0);

  IMAGE_IMPORT_DESCRIPTOR *ImportDescriptor = FindImportDescriptor(DllName);
  if(ImportDescriptor == NULL) return FALSE;

  // Recreate orginal imports
  IMAGE_THUNK_DATA32 *FThunk, *OThunk;
  FThunk = (IMAGE_THUNK_DATA32*)(MainVA + ImportDescriptor->FirstThunk);
  OThunk = (IMAGE_THUNK_DATA32*)(MainVA + ImportDescriptor->OriginalFirstThunk);

//  HINSTANCE OriginalDll = LoadLibrary(OriginalDllPath);
//  if(OriginalDll == NULL) return FALSE;

  for(; OThunk->u1.AddressOfData; OThunk++, FThunk++)
  {
    DWORD OldRight;
    VirtualProtect(&FThunk->u1.Function, 4, PAGE_EXECUTE_READWRITE, &OldRight);

    if(OThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG32)
    {
      FThunk->u1.Function = (DWORD)GetProcAddress(hOriginal, (const char*)(DWORD)(WORD)OThunk->u1.Ordinal);;
    }
    else
    {
      IMAGE_IMPORT_BY_NAME *ImportFuncName = (IMAGE_IMPORT_BY_NAME*)(MainVA + OThunk->u1.AddressOfData);
      FThunk->u1.Function = (DWORD)GetProcAddress(hOriginal, (const char*)ImportFuncName->Name);
    }

    DWORD Temp;
    VirtualProtect(&FThunk->u1.Function, 4, OldRight, &Temp);
  }

  return TRUE;
}

static IMAGE_EXPORT_DIRECTORY* FindExportDirectory(HINSTANCE hInst)
{
  DWORD InstVA = (DWORD)hInst;

  IMAGE_DOS_HEADER *DosHeader = (IMAGE_DOS_HEADER*)hInst;
  IMAGE_NT_HEADERS *NtHeaders = (IMAGE_NT_HEADERS*)(DosHeader->e_lfanew + InstVA);
  IMAGE_EXPORT_DIRECTORY* ExpDir = 
    (IMAGE_EXPORT_DIRECTORY*)(NtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress + InstVA);
  
  return ExpDir;
}

BOOL DllSpoof_HackExports(HINSTANCE hSelf, HINSTANCE hOriginal)
{
  DWORD SelfVA = (DWORD)hSelf;
  IMAGE_EXPORT_DIRECTORY *SelfExp = FindExportDirectory(hSelf);

  int i;
  DWORD *SelfNames = (DWORD*)(SelfVA + SelfExp->AddressOfNames);
  DWORD *SelfFuncAddr = (DWORD*)(SelfVA + SelfExp->AddressOfFunctions);

  DWORD OldRight;
  VirtualProtect(SelfFuncAddr, 4 * SelfExp->NumberOfNames, PAGE_EXECUTE_READWRITE, &OldRight);

  for(i = 0; i < SelfExp->NumberOfNames; i++)
  {
    const char *SelfName = (const char*)(SelfVA + SelfNames[i]);
    if((DWORD)SelfName == SelfVA) continue;
    SelfFuncAddr[i] = (DWORD)GetProcAddress(hOriginal, SelfName) - SelfVA;

    //printf("%s == %x\n", SelfName, SelfFuncAddr[i]);
    //fflush(stdout);
  }

  DWORD Temp;
  VirtualProtect(SelfFuncAddr, 4 * SelfExp->NumberOfNames, OldRight, &Temp);
  return TRUE;
}

